import requests
import json

api_url = "http://10.1.161.53:8080/api/v1/chats/28792a821a6e11f090350242ac180006/completions"
api_url_sessions = "http://10.1.161.53:8080/api/v1/chats/28792a821a6e11f090350242ac180006/sessions"
api_key = "ragflow-AzNDhiZWNlMWE2YjExZjA5ZjlkMDI0Mm"  # 替换为实际API密钥


def create_session(question):
    """创建新会话并返回session_id"""
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }
    payload = {"name": question}

    try:
        response = requests.post(api_url_sessions, headers=headers, json=payload)
        response.raise_for_status()
        return response.json()["data"]["id"]
    except requests.exceptions.RequestException as e:
        print(f"创建会话失败: {e}")
        return None


def stream_question_generator(question, session_id):
    """生成器方法，返回实时流片段"""
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }
    payload = {
        "question": question,
        "stream": True,
        "session_id": session_id
    }

    prev_answer = ""
    buffer = ""

    try:
        with requests.post(api_url, headers=headers, json=payload, stream=True) as response:
            response.raise_for_status()
            for chunk in response.iter_content(chunk_size=8192, decode_unicode=True):
                if chunk:
                    buffer += chunk
                    while '\n' in buffer:
                        line, buffer = buffer.split('\n', 1)
                        line = line.strip()
                        if line.startswith('data:'):
                            try:
                                # 处理结束标记
                                if line == 'data:{"code": 0, "data": true}':
                                    return  # 结束生成器

                                # 处理正常数据
                                data = json.loads(line[5:])
                                current_answer = data.get('data', {}).get('answer', '')

                                # 计算新增部分
                                new_part = current_answer[len(prev_answer):]
                                if new_part:
                                    prev_answer = current_answer
                                    yield new_part  # 返回新增片段

                            except json.JSONDecodeError:
                                continue  # 跳过无法解析的块
    except requests.exceptions.RequestException as e:
        yield f"请求发生错误: {e}"


# 使用示例
# if __name__ == "__main__":
#     QUESTION = "关于救生衣的最新政策"
#     session_id = create_session(QUESTION)
#     print(f"创建会话ID: {session_id}")
#
#     for fragment in stream_question_generator(QUESTION, session_id):
#         print(fragment, end='', flush=True)
#     print()